from utils import load_yaml_file
from data_source import get_data_by_meta
from data_process import TjdSingleData

import pandas as pd

data_sets_config = load_yaml_file('pred.yaml')


def extract_series_from_df(df, name):
    if name in df.columns:
        return 1, df[name]
    else:
        return 0, pd.Series()


def max_to_third(se):
    if max(se) / abs(se.mean()) > 20:
        print(se.name, '极端值修正')
        max_v = se.max()
        max_3 = max_v / 3
        se.loc[se == max_v] = max_3
    return se


def prepare_datasets_dic(y_name, meta_data):
    # 定义一个列表，包含两个字符串，分别为'通信数据'和'油气数据'
    alter_names = ['通信数据', '油气数据']
    # 使用字典推导式，从'./raw_data/'目录下读取两个xlsx文件，并将它们存储在字典alternate_data中
    alternate_data = {s : pd.read_excel(f'./raw_data/{s}.xlsx', index_col=0) for s in alter_names}

    # 从data_sets_config字典中获取与y_name对应的x_name_dic
    x_name_dic = data_sets_config['data_set'][y_name]
    # 定义一个空字典x_data_dic，用于存储处理后的x数据
    x_data_dic = {}
    # 遍历x_name_dic中的每一个键值对
    for x_name, (x_name_process, x_source) in x_name_dic.items():
        # 如果x_source为'meta'，则从meta_data中提取x_name对应的序列
        if x_source == 'meta':
            x_get, x_se = extract_series_from_df(meta_data, x_name)
        # 否则，从alternate_data中提取x_name对应的序列
        else:
            df_alter = alternate_data[x_source]
            x_get, x_se = extract_series_from_df(df_alter, x_name)
        # 如果x_get为True，则进行后续处理
        if x_get:
            
            # 删除x_se中的缺失值
            x_se = x_se.dropna()
            # 使用TjdSingleData类对x_se进行处理，得到df_hf
            df_hf = TjdSingleData(x_se)
            
            # 如果x_name_process的长度小于2，则直接将df_hf.data赋值给df_x
            if len(x_name_process) < 2:
                df_x = df_hf.data
            # 否则，根据x_name_process的值，从df_hf中提取对应的数据，赋值给df_x
            else:
                if '_' in x_name_process:
                    df_x = df_hf.get_long_string_data(x_name_process).data
                else:
                    df_x = df_hf[x_name_process].data
            # 打印y_name，'feature'，x_name，x_name_process
            print(y_name, '>>', 'feature', '>>', x_name, ' process ', x_name_process)
            x_data_dic[x_name] = max_to_third(df_x)
            x_data_dic[x_name] = df_x
    x_data_dic['工业增加值当月同比']['2025-03-31'] = 5.9
    return x_data_dic


def prepare_xy_data(y_name: str) -> tuple[pd.DataFrame]:
    
    # 加载配置文件
    meta = load_yaml_file('config.yaml')
    # 根据配置文件中的meta字段获取数据
    meta_data = get_data_by_meta(meta['meta'])
    # 根据y_name字段准备数据集字典
    x_datasets = prepare_datasets_dic(y_name, meta_data)
    # 根据y_name字段获取y_data数据，并删除缺失值
    y_data = meta_data.loc[:,[y_name]].dropna()
    # 将y_data的索引转换为日期时间格式
    y_data.index = pd.to_datetime(y_data.index)
    # 返回y_data和x_datasets
    return y_data, x_datasets